from tracemalloc import start
from typing import List, Optional
from datetime import datetime, timedelta
import pandas as pd
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import BarData, HistoryRequest
from vnpy.trader.database import get_database
from vnpy.trader.datafeed import BaseDatafeed
from vnpy.trader.setting import SETTINGS
from vnpy_tushare.tushare_datafeed import TushareDatafeed
from portfolio_trader.datafeed.storage import DataStorage
import os


class PortfolioDataLoader:
    """
    投资组合数据加载器
    """
    def __init__(self):
        """初始化"""
        self.database = get_database()  # 数据库
        self.datafeed: BaseDatafeed = TushareDatafeed()  # Tushare数据源
        self.pro_api = None
        self.storage = DataStorage()
        
    def init(self, tushare_token: str) -> bool:
        """初始化Tushare token"""
        SETTINGS['datafeed.username'] = "token"
        SETTINGS["datafeed.password"] = tushare_token
        if self.datafeed.init():
            import tushare as ts
            ts.set_token(tushare_token)
            self.pro_api = ts.pro_api()
            return True
        return False

    def _get_exchange(self, symbol: str) -> Exchange:
        """获取交易所（Exchange.SSE/Exchange.SZSE）"""
        if symbol.startswith(("600", "601", "603", "688")):
            return Exchange.SSE
        elif symbol.startswith(("000", "002", "300")):
            return Exchange.SZSE
        elif symbol.startswith(("4", "8")):
            return Exchange.BSE
        else:
            return Exchange.SSE  # 默认返回上交所 

    def _get_exchange_code(self, symbol: str) -> str:
        """获取交易所代码（SH/SZ）"""
        exchange = self._get_exchange(symbol)
        if exchange == Exchange.SSE:
            return "SH"
        elif exchange == Exchange.SZSE:
            return "SZ"
        else:
            return ""

    def get_stock_name(self, symbol: str) -> str:
        """获取股票名称"""
        # 通过pro_api获取
        if self.pro_api:
            try:
                ts_code = f"{symbol}.{self._get_exchange_code(symbol)}"
                df = self.pro_api.namechange(
                    ts_code=ts_code,
                    fields="ts_code,name"
                )
                if df is not None and not df.empty:
                    # 返回第一行的name列，即为当前股票名称
                    return df.iloc[0]["name"]
            except Exception as e:
                print(f"通过pro_api获取股票名称失败: {str(e)}")
        else:
            print("Tushare API 未初始化")
            
        return ""

    def get_latest_local_date(self, symbol: str, data_type: str = "daily") -> Optional[pd.Timestamp]:
        """
        获取本地数据的最新日期
        """
        if self.storage.has_data(symbol, data_type):
            csv_path = os.path.join(self.storage.data_path, f"{symbol}_{data_type}.csv")
            df = pd.read_csv(csv_path)
            if not df.empty:
                df["datetime"] = pd.to_datetime(df["datetime"])
                return df["datetime"].max()
        else:
            print(f"本地没有 {symbol} 的 {data_type} 数据")
        return None

    def get_trade_calendar(self, symbol: str, start_date: datetime, end_date: datetime) -> pd.DataFrame:
        """获取交易日历
        :param symbol: 股票代码
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: DataFrame，范围为[start_date, end_date]内的交易日期
        """
        if not self.pro_api:
            print("Tushare API 未初始化 from get_trade_calendar")
            return pd.DataFrame()
        
        # 获取交易所代码
        exchange = self._get_exchange(symbol)
        
        # 如果当前时间早于收盘时间(15:00)，则排除当天日期
        current_time = datetime.now().replace(tzinfo=None)
        if current_time.hour < 15 and end_date.replace(tzinfo=None).date() == current_time.date():
            print("当前时间早于15:00，排除当天交易日历")
            end_date = end_date - timedelta(days=1)
            if end_date.replace(tzinfo=None) < start_date.replace(tzinfo=None):
                print("调整后的结束日期早于开始日期，返回空DataFrame")
                return pd.DataFrame()
        
        start_date = start_date.strftime("%Y%m%d")
        end_date = end_date.strftime("%Y%m%d")

        # 获取完整的交易日历信息
        calendar = self.pro_api.trade_cal(
            exchange=exchange.value,
            start_date=start_date,
            end_date=end_date,
            is_open=1,
            fields="cal_date"
        )
        print(f"交易日历开始时间: {start_date}, 结束时间: {end_date}")
        if calendar is not None and not calendar.empty:
            # 转换日期格式
            calendar["cal_date"] = pd.to_datetime(calendar["cal_date"])
            
            # 按日期降序排序
            calendar = calendar.sort_values("cal_date", ascending=False)

            print(f"获取到交易日历，开始时间: {calendar.iloc[-1]['cal_date']}, 结束时间: {calendar.iloc[0]['cal_date']}")
            
            return calendar
        else:
            print(f"未获取到交易日历数据")
            return pd.DataFrame()

    def _get_daily_data(self, symbol: str, start_date: datetime, end_date: datetime) -> List[BarData]:
        """获取日K线数据
        :param symbol: 股票代码
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: BarData列表
        """
        exchange = self._get_exchange(symbol)
        req = HistoryRequest(
            symbol=symbol,
            exchange=exchange,
            start=start_date,
            end=end_date,
            interval=Interval.DAILY
        )
        bars = self.datafeed.query_bar_history(req)
        if bars:
            print(f"从API获取到 {len(bars)} 条日K线数据")
            self.storage.save_bar_data(bars, symbol)
            return bars
        else:
            print("从API未获取到数据")
            return []

    def _get_valuation_data(self, symbol: str, start_date: datetime, end_date: datetime) -> pd.DataFrame:
        """获取估值数据
        :param symbol: 股票代码
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: DataFrame
        """
        exchange_code = self._get_exchange_code(symbol)
        ts_code = f"{symbol}.{exchange_code}"
        
        valuation = self.pro_api.daily_basic(
            ts_code=ts_code,
            start_date=start_date.strftime("%Y%m%d"),
            end_date=end_date.strftime("%Y%m%d"),
            fields="ts_code,trade_date,pe,pb,ps,dv_ratio,free_share,circ_mv"
        )
        
        if valuation is not None and not valuation.empty:
            print(f"从API获取到 {len(valuation)} 条估值数据")
            valuation = valuation.rename(columns={"ts_code": "symbol", "trade_date": "datetime"})
            valuation["exchange"] = valuation["symbol"].str.split(".").str[1].map({"SH": "SSE", "SZ": "SZSE"})
            valuation['symbol'] = valuation['symbol'].str.split('.').str[0]
            cols = valuation.columns.tolist()
            cols = cols[:1] + ['exchange'] + cols[1:-1]
            valuation = valuation[cols]
            valuation["datetime"] = pd.to_datetime(valuation["datetime"])
            self.storage.save_valuation_data(valuation, symbol)
            return valuation
        else:
            print("从API未获取到估值数据")
            return pd.DataFrame()

    def update_daily_data(self, symbol: str, end_date: datetime) -> None:
        """更新本地数据
        :param symbol: 股票代码
        :param end_date: 结束日期
        """
        latest_date = self.get_latest_local_date(symbol)
        if not latest_date:
            # 如果本地没有数据，获取近一年的数据
            start_date = end_date - timedelta(days=365)
            self._get_daily_data(symbol, start_date, end_date)
            return

        # 获取交易日历
        calendar = self.get_trade_calendar(symbol, latest_date, end_date)
        if calendar.empty or len(calendar) <= 1:  # 排除最后一天
            return

        # 发现新的交易日数据，从API更新
        print(f"发现新的交易日数据，从API更新")
        self._get_daily_data(symbol, latest_date, end_date)

    def update_valuation_data(self, symbol: str, end_date: datetime) -> None:
        """更新本地估值数据
        :param symbol: 股票代码
        :param end_date: 结束日期
        """
        latest_date = self.get_latest_local_date(symbol, "valuation")
        if not latest_date:
            # 如果本地没有数据，获取近一年的数据
            start_date = end_date - timedelta(days=365)
            self._get_valuation_data(symbol, start_date, end_date)
            return

        # 获取交易日历
        calendar = self.get_trade_calendar(symbol, latest_date, end_date)
        if calendar.empty or len(calendar) <= 1:  # 排除最后一天
            return

        # 发现新的交易日数据，从API更新估值数据
        print(f"发现新的交易日数据，从API更新估值数据")
        self._get_valuation_data(symbol, latest_date, end_date)

    def load_daily_data(
        self,
        symbol: str,
        start_date: datetime,
        end_date: datetime
    ) -> List[BarData]:
        """
        加载日线数据
        :param symbol: 股票代码
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: BarData列表
        """
        exchange = self._get_exchange(symbol)
        
        # 先更新本地数据
        self.update_daily_data(symbol, end_date)
        
        # 从本地加载指定时间范围的数据
        if self.storage.has_data(symbol, "daily"):
            csv_path = os.path.join(self.storage.data_path, f"{symbol}_daily.csv")
            df = pd.read_csv(csv_path)
            df["datetime"] = pd.to_datetime(df["datetime"]).dt.tz_localize(None)
            
            # 过滤时间范围并按时间降序排序
            df_date = df["datetime"].dt.date
            start_date_only = start_date.date()
            end_date_only = end_date.date()
            df = df[(df_date >= start_date_only) & (df_date <= end_date_only)].sort_values("datetime", ascending=False)
            
            if not df.empty:
                bars = []
                for _, row in df.iterrows():
                    bar = BarData(
                        gateway_name="Tushare",
                        symbol=row["symbol"],
                        exchange=Exchange(row["exchange"]),
                        datetime=row["datetime"].replace(hour=15),
                        interval=Interval(row["interval"]),
                        volume=row["volume"],
                        open_price=row["open_price"],
                        high_price=row["high_price"],
                        low_price=row["low_price"],
                        close_price=row["close_price"]
                    )
                    bars.append(bar)
                print(f"从本地加载了 {len(bars)} 条数据")
                return bars
        
        # 如果本地没有数据，则从API获取
        print(f"从API获取股票 {symbol} 的日线数据，时间范围: {start_date} 到 {end_date}")
        bars = self._get_daily_data(symbol, start_date, end_date)
        if bars:
            return bars
        else:
            return []

    def load_valuation_data(
        self,
        symbol: str,
        start_date: datetime,
        end_date: datetime
    ) -> pd.DataFrame:
        """
        加载估值数据
        :param symbol: 股票代码
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: DataFrame
        """
        if not self.pro_api:
            print("Tushare API 未初始化")
            return pd.DataFrame()
        
        # 先更新本地数据
        self.update_valuation_data(symbol, end_date)
        
        # 从本地加载指定时间范围的数据
        if self.storage.has_data(symbol, "valuation"):
            csv_path = os.path.join(self.storage.data_path, f"{symbol}_valuation.csv")
            valuation = pd.read_csv(csv_path)
            valuation["datetime"] = pd.to_datetime(valuation["datetime"])
            
            # 过滤时间范围
            valuation = valuation[(valuation["datetime"] >= start_date) & (valuation["datetime"] <= end_date)]
            
            if not valuation.empty:
                print(f"从本地加载了 {len(valuation)} 条估值数据")
                return valuation
        
        # 如果本地没有数据，则从API获取
        print(f"从API获取 {symbol} 的估值数据，时间范围: {start_date} 到 {end_date}")
        
        valuation = self._get_valuation_data(symbol, start_date, end_date)
        
        if valuation is not None and not valuation.empty:
            return valuation
        else:
            return pd.DataFrame()

    def get_latest_daily_data(self, symbol: str) -> Optional[BarData]:
        """获取最新的日K线数据
        :param symbol: 股票代码
        :return: 最新的K线数据，如果没有数据则返回None
        """
        current_time = datetime.now()
        # 先更新数据
        self.update_daily_data(symbol, current_time)
        
        # 从本地加载最新数据
        if self.storage.has_data(symbol, "daily"):
            csv_path = os.path.join(self.storage.data_path, f"{symbol}_daily.csv")
            df = pd.read_csv(csv_path)
            if not df.empty:
                df["datetime"] = pd.to_datetime(df["datetime"])
                # 获取最新的一条记录
                latest_row = df.sort_values("datetime", ascending=False).iloc[0]
                print(f"从本地加载了最新的 {symbol} 日K线数据")
                return BarData(
                    gateway_name="Tushare",
                    symbol=latest_row["symbol"],
                    exchange=Exchange(latest_row["exchange"]),
                    datetime=latest_row["datetime"].replace(hour=15),
                    interval=Interval(latest_row["interval"]),
                    volume=latest_row["volume"],
                    open_price=latest_row["open_price"],
                    high_price=latest_row["high_price"],
                    low_price=latest_row["low_price"],
                    close_price=latest_row["close_price"]
                )
        print("本地数据加载失败")
        return None

    def get_latest_valuation_data(self, symbol: str) -> Optional[pd.Series]:
        """获取最新的估值数据
        :param symbol: 股票代码
        :return: 最新的估值数据，如果没有数据则返回None
        """
        current_time = datetime.now()
        # 先更新数据
        self.update_valuation_data(symbol, current_time)
        
        # 从本地加载最新数据
        if self.storage.has_data(symbol, "valuation"):
            csv_path = os.path.join(self.storage.data_path, f"{symbol}_valuation.csv")
            df = pd.read_csv(csv_path)
            if not df.empty:
                df["datetime"] = pd.to_datetime(df["datetime"])
                # 获取最新的一条记录
                latest_row = df.sort_values("datetime", ascending=False).iloc[0]
                print(f"从本地加载了最新的 {symbol} 估值数据")
                return latest_row
        print("本地最新估值数据加载失败")
        return None
